#include <ros/ros.h>
#include <opencv2/opencv.hpp>
#include <opencv2/calib3d.hpp>
#include <fstream>
#include <sstream>
#include <nav_msgs/Odometry.h>
#include <nav_msgs/Path.h>
#include <tf/transform_broadcaster.h>
#include <vector>
#include <string>
#include <map>

// 读取相机内参的函数
bool readCameraParameters(const std::string& filename, cv::Mat& camera_matrix, cv::Mat& dist_coeffs)
{
    cv::FileStorage fs(filename, cv::FileStorage::READ);
    if (!fs.isOpened())
    {
        ROS_ERROR("Failed to open camera parameter file: %s", filename.c_str());
        return false;
    }

    fs["projection_parameters"]["fx"] >> camera_matrix.at<double>(0, 0);
    fs["projection_parameters"]["fy"] >> camera_matrix.at<double>(1, 1);
    fs["projection_parameters"]["cx"] >> camera_matrix.at<double>(0, 2);
    fs["projection_parameters"]["cy"] >> camera_matrix.at<double>(1, 2);

    camera_matrix.at<double>(0, 1) = 0; // skew is zero
    camera_matrix.at<double>(1, 0) = 0;
    camera_matrix.at<double>(2, 0) = 0;
    camera_matrix.at<double>(2, 1) = 0;
    camera_matrix.at<double>(2, 2) = 1;

    fs["distortion_parameters"]["k1"] >> dist_coeffs.at<double>(0);
    fs["distortion_parameters"]["k2"] >> dist_coeffs.at<double>(1);
    fs["distortion_parameters"]["p1"] >> dist_coeffs.at<double>(2);
    fs["distortion_parameters"]["p2"] >> dist_coeffs.at<double>(3);

    return true;
}

class PnPNode
{
public:
    PnPNode(const std::string& points_filename, const std::string& camera_filename,const std::string& output_filename, ros::NodeHandle& nh) : nh_(nh), current_image_index_(0)
    {
        match_path_.header.frame_id = "world";

        if (!readCameraParameters(camera_filename, camera_matrix, dist_coeffs))
        {
            ROS_ERROR("Failed to read camera parameters.");
            return;
        }
        if (!readPointsFromFile(points_filename))
        {
            ROS_ERROR("Failed to read points from file.");
            return;
        }

        odometry_pub_ = nh_.advertise<nav_msgs::Odometry>("pnp_odometry", 10);
        path_pub_ = nh_.advertise<nav_msgs::Path>("pnp_path", 10);
        // output_file_.open("/home/luli/pnp_ws/src/pnp_utils/match_result/playground_200_incline.csv");
        output_file_.open(output_filename);

        if (!output_file_.is_open())
        {
            ROS_ERROR("Failed to open output CSV file.");
        }
        else
        {
            output_file_ << "Timestamp,ImageName,X,Y,Z\n";
        }

        timer_ = nh_.createTimer(ros::Duration(0.1), &PnPNode::processNextImage, this);
    }

    ~PnPNode()
    {
        if (output_file_.is_open())
        {
            output_file_.close();
        }
    }

private:
    ros::NodeHandle nh_;
    ros::Publisher odometry_pub_, path_pub_;
    ros::Timer timer_;
    cv::Mat camera_matrix = cv::Mat::zeros(3, 3, CV_64F);
    cv::Mat dist_coeffs = cv::Mat::zeros(4, 1, CV_64F);
    std::map<std::string, std::vector<cv::Point2f>> image_points_map;
    std::map<std::string, std::vector<cv::Point3f>> object_points_map;
    std::map<std::string, double> timestamps_map;  // Store timestamps for each image
    std::ofstream output_file_;
    size_t current_image_index_;
    nav_msgs::Path match_path_;

    bool readPointsFromFile(const std::string& filename)
    {
        std::ifstream file(filename);
        if (!file.is_open())
        {
            ROS_ERROR("Failed to open file: %s", filename.c_str());
            return false;
        }

        std::string line;
        while (std::getline(file, line))
        {
            std::istringstream iss(line);
            std::string token;
            std::vector<std::string> tokens;

            while (std::getline(iss, token, ','))
            {
                tokens.push_back(token);
            }

            if (tokens.size() != 8)
            {
                ROS_ERROR("Failed to parse line: %s", line.c_str());
                continue;
            }

            float x_pixel = std::stof(tokens[0]);
            float y_pixel = std::stof(tokens[1]);
            float x_3D = std::stof(tokens[2]);
            float y_3D = std::stof(tokens[3]);
            float z_3D = std::stof(tokens[4]);
            std::string image_name = tokens[5];
            double timestamp = std::stod(tokens[6]);
            int frame_id = std::stoi(tokens[7]);

            image_points_map[image_name].push_back(cv::Point2f(x_pixel, y_pixel));
            object_points_map[image_name].push_back(cv::Point3f(x_3D, y_3D, z_3D));
            timestamps_map[image_name] = timestamp;  // Store the timestamp
        }
        file.close();

        return true;
    }

    void processNextImage(const ros::TimerEvent&)
    {
        if (current_image_index_ >= image_points_map.size())
        {
            ROS_INFO_STREAM_THROTTLE(15.0, "\033[35m All images processed.   \033[0m");
            return;
        }

        auto it = image_points_map.begin();
        std::advance(it, current_image_index_);
        const std::string& image_name = it->first;
        const std::vector<cv::Point2f>& image_points = it->second;
        const std::vector<cv::Point3f>& object_points = object_points_map[image_name];
        double timestamp = timestamps_map[image_name];  // Get the timestamp

        if (image_points.size() < 5)
        {
            ROS_WARN("Insufficient points for image: %s", image_name.c_str());
            current_image_index_++;
            return;
        }

        cv::Mat rvec, tvec;
        bool success = cv::solvePnP(object_points, image_points, camera_matrix, dist_coeffs, rvec, tvec, false, cv::SOLVEPNP_EPNP);

        if (success && abs(tvec.at<double>(0))<= 10000)
        {
            cv::Mat rotation_matrix;
            cv::Rodrigues(rvec, rotation_matrix);
            ROS_INFO_STREAM("\033[32m Translation Vector for " << image_name << ": \n \033[0m" << tvec);

            // Publish odometry
            nav_msgs::Odometry odometry_msg;
            
            odometry_msg.header.stamp = ros::Time().fromNSec(timestamp );  // Use the timestamp from the file
            // odometry_msg.header.stamp = ros::Time().fromNSec(timestamp );  // Use the timestamp from the file
            odometry_msg.header.frame_id = "world";
            odometry_msg.pose.pose.position.x = tvec.at<double>(0);
            odometry_msg.pose.pose.position.y = tvec.at<double>(1);
            odometry_msg.pose.pose.position.z = tvec.at<double>(2);

            tf::Quaternion q;
            tf::Matrix3x3(
                rotation_matrix.at<double>(0, 0), rotation_matrix.at<double>(0, 1), rotation_matrix.at<double>(0, 2),
                rotation_matrix.at<double>(1, 0), rotation_matrix.at<double>(1, 1), rotation_matrix.at<double>(1, 2),
                rotation_matrix.at<double>(2, 0), rotation_matrix.at<double>(2, 1), rotation_matrix.at<double>(2, 2)
            ).getRotation(q);

            odometry_msg.pose.pose.orientation.x = q.x();
            odometry_msg.pose.pose.orientation.y = q.y();
            odometry_msg.pose.pose.orientation.z = q.z();
            odometry_msg.pose.pose.orientation.w = q.w();

            odometry_pub_.publish(odometry_msg);

            // Publish path
            geometry_msgs::PoseStamped pose_stamped;
            pose_stamped.header = odometry_msg.header;
            pose_stamped.pose = odometry_msg.pose.pose;
            match_path_.poses.push_back(pose_stamped);

            path_pub_.publish(match_path_);

            std::string time_s ;
            time_s = std::to_string(timestamp);

            if (output_file_.is_open())
            {
                output_file_ << time_s << "," << image_name << ","
                             << tvec.at<double>(0) << "," << tvec.at<double>(1) << "," << tvec.at<double>(2) << "\n";
            }
        }
        else
        {
            ROS_ERROR("solvePnP failed to find a solution for image: %s", image_name.c_str());
        }

        current_image_index_++;
    }
};

int main(int argc, char** argv)
{
    ros::init(argc, argv, "pnp_node");
    ros::NodeHandle nh;

    std::string points_filename = "/home/luli/pnp_ws/match_result/output.txt";
    std::string camera_filename = "/home/luli/pnp_ws/src/pnp_utils/config/downcam.yaml";
    std::string output_filename = "/home/luli/pnp_ws/src/pnp_utils/config/result.yaml";

    nh.getParam("/points_filename", points_filename);
    nh.getParam("/camera_filename", camera_filename);
    nh.getParam("/output_filename", output_filename);

    std::cout << "points_filename : " << points_filename << std::endl;
    std::cout << "camera_filename : " << camera_filename << std::endl;
    std::cout << "output_filename : " << output_filename << std::endl;

    PnPNode node(points_filename, camera_filename,output_filename, nh);

    ros::spin();
    return 0;
}
